import os
import asyncio
import finnhub
import pandas as pd
import json
import random
from typing import Annotated
from collections import defaultdict
from fastmcp import FastMCP, Client
from datetime import datetime
from dotenv import load_dotenv

# 加载 .env 文件中的环境变量
load_dotenv()

# 创建FastMCP客户端实例
mcp = FastMCP('demo.mcp')

# 初始化 Finnhub 客户端
def init_finnhub_client():
    # 从环境变量中获取 API 密钥
    finnhub_api_key = os.getenv("FINNHUB_API_KEY")
    
    if finnhub_api_key is None:
        print("Please set the environment variable FINNHUB_API_KEY to use the Finnhub API.")
        return None
    else:
        return finnhub.Client(api_key=finnhub_api_key)

# 公司的基本财务数据
@mcp.tool()
def get_basic_financials(symbol: Annotated[str, "ticker symbol"], selected_columns: Annotated[list[str], "List of selected columns"]) -> str:
    """
    获取指定公司的最新财务数据
    :param symbol: 公司的股票符号
    :param selected_columns: 要选择的财务数据列
    :return: 返回所选择财务数据的JSON格式字符串
    """
    finnhub_client = init_finnhub_client()
    if not finnhub_client:
        return "Finnhub client initialization failed!"
    
    basic_financials = finnhub_client.company_basic_financials(symbol, "all")
    if not basic_financials["series"]:
        return f"Failed to find basic financials for symbol {symbol} from finnhub! Try a different symbol."
    
    output_dict = basic_financials["metric"]
    for metric, value_list in basic_financials["series"]["quarterly"].items():
        value = value_list[0]
        output_dict.update({metric: value["v"]})

    # 筛选所需的列
    for k in list(output_dict.keys()):
        if selected_columns and k not in selected_columns:
            output_dict.pop(k)
    
    return str(output_dict)

# 公司的简介
@mcp.tool()
def get_company_profile(symbol: Annotated[str, "ticker symbol"]) -> str:
    """
    获取指定公司的简介
    """
    finnhub_client = init_finnhub_client()
    if not finnhub_client:
        return "Finnhub client initialization failed!"

    profile = finnhub_client.company_profile2(symbol=symbol)
    if not profile:
        return f"Failed to find company profile for symbol {symbol} from finnhub!"

    formatted_str = (
        "[Company Introduction]:\n\n{name} is a leading entity in the {finnhubIndustry} sector. "
        "Incorporated and publicly traded since {ipo}, the company has established its reputation as "
        "one of the key players in the market. As of today, {name} has a market capitalization "
        "of {marketCapitalization:.2f} in {currency}, with {shareOutstanding:.2f} shares outstanding."
        "\n\n{name} operates primarily in the {country}, trading under the ticker {ticker} on the {exchange}. "
        "As a dominant force in the {finnhubIndustry} space, the company continues to innovate and drive "
        "progress within the industry."
    ).format(**profile)

    return formatted_str

# 获取公司新闻
@mcp.tool()
def get_company_news(
    symbol: Annotated[str, "ticker symbol"],
    start_date: Annotated[str, "start date of the search period for the company's basic financials, yyyy-mm-dd"],
    end_date: Annotated[str, "end date of the search period for the company's basic financials, yyyy-mm-dd"],
    max_news_num: Annotated[int, "maximum number of news to return, default to 10"] = 10
) -> pd.DataFrame:
    """
    获取指定公司在指定日期范围内的新闻
    """
    finnhub_client = init_finnhub_client()
    if not finnhub_client:
        return "Finnhub client initialization failed!"
    
    news = finnhub_client.company_news(symbol, _from=start_date, to=end_date)
    if len(news) == 0:
        return f"No company news found for symbol {symbol} from finnhub!"
    
    news = [
        {
            "date": datetime.fromtimestamp(n["datetime"]).strftime("%Y%m%d%H%M%S"),
            "headline": n["headline"],
            "summary": n["summary"],
        }
        for n in news
    ]
    
    # 随机选择新闻
    if len(news) > max_news_num:
        news = random.choices(news, k=max_news_num)
    
    news.sort(key=lambda x: x["date"])
    output = pd.DataFrame(news)

    return output

# 获取公司财务数据历史记录
@mcp.tool()
def get_basic_financials_history(
    symbol: Annotated[str, "ticker symbol"],
    freq: Annotated[str, "reporting frequency of the company's basic financials: annual / quarterly"],
    start_date: Annotated[str, "start date of the search period for the company's basic financials, yyyy-mm-dd"],
    end_date: Annotated[str, "end date of the search period for the company's basic financials, yyyy-mm-dd"],
    selected_columns: Annotated[list[str] | None, "List of selected columns"]
) -> pd.DataFrame:
    """
    获取公司指定财务数据历史记录
    """
    finnhub_client = init_finnhub_client()
    if not finnhub_client:
        return "Finnhub client initialization failed!"
    
    if freq not in ["annual", "quarterly"]:
        return f"Invalid reporting frequency {freq}. Please specify either 'annual' or 'quarterly'."

    basic_financials = finnhub_client.company_basic_financials(symbol, "all")
    if not basic_financials["series"]:
        return f"Failed to find basic financials for symbol {symbol} from finnhub! Try a different symbol."

    output_dict = defaultdict(dict)
    for metric, value_list in basic_financials["series"][freq].items():
        if selected_columns and metric not in selected_columns:
            continue
        for value in value_list:
            if value["period"] >= start_date and value["period"] <= end_date:
                output_dict[metric].update({value["period"]: value["v"]})

    financials_output = pd.DataFrame(output_dict)
    financials_output = financials_output.rename_axis(index="date")

    return financials_output

GET_BASIC_FINANCIALS_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_basic_financials",
        "description": "获取指定公司的最新财务数据，并按所给列进行筛选。",
        "parameters": {
            "type": "object",
            "properties": {
                "symbol": {
                    "type": "string",
                    "description": "股票代码（如 'AAPL'）。"
                },
                "selected_columns": {
                    "type": "array",
                    "description": "需要返回的财务字段列表；空或不传表示返回全部。",
                    "items": {"type": "string"},
                    "default": []
                }
            },
            "required": ["symbol"]
        }
    }
}

GET_COMPANY_PROFILE_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_company_profile",
        "description": "获取指定公司的公司简介（交易所、行业、市值等）。",
        "parameters": {
            "type": "object",
            "properties": {
                "symbol": {
                    "type": "string",
                    "description": "股票代码（如 'AAPL'）。"
                }
            },
            "required": ["symbol"]
        }
    }
}

GET_COMPANY_NEWS_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_company_news",
        "description": "获取公司在指定日期区间内的新闻（按时间排序，可限制条数）。",
        "parameters": {
            "type": "object",
            "properties": {
                "symbol": {
                    "type": "string",
                    "description": "股票代码（如 'AAPL'）。"
                },
                "start_date": {
                    "type": "string",
                    "description": "开始日期（YYYY-MM-DD）。",
                    "format": "date"
                },
                "end_date": {
                    "type": "string",
                    "description": "结束日期（YYYY-MM-DD）。",
                    "format": "date"
                },
                "max_news_num": {
                    "type": "integer",
                    "description": "返回的新闻最大条数。",
                    "minimum": 1,
                    "maximum": 100,
                    "default": 10
                }
            },
            "required": ["symbol", "start_date", "end_date"]
        }
    }
}

GET_BASIC_FINANCIALS_HISTORY_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_basic_financials_history",
        "description": "按频率（annual/quarterly）获取公司指定财务字段在日期区间内的历史记录。",
        "parameters": {
            "type": "object",
            "properties": {
                "symbol": {
                    "type": "string",
                    "description": "股票代码（如 'AAPL'）。"
                },
                "freq": {
                    "type": "string",
                    "description": "报表频率：'annual' 或 'quarterly'。",
                    "enum": ["annual", "quarterly"]
                },
                "start_date": {
                    "type": "string",
                    "description": "开始日期（YYYY-MM-DD）。",
                    "format": "date"
                },
                "end_date": {
                    "type": "string",
                    "description": "结束日期（YYYY-MM-DD）。",
                    "format": "date"
                },
                "selected_columns": {
                    "type": "array",
                    "description": "需要返回的财务字段列表；空或不传表示返回全部。",
                    "items": {"type": "string"},
                    "default": []
                }
            },
            "required": ["symbol", "freq", "start_date", "end_date"]
        }
    }
}


# 异步调用主函数
async def main():
    client = Client(mcp)
    async with client:
        # 查看可用工具
        tools = await client.list_tools()
        print('可用工具:', tools, "\n")  
        print("-------------------------------------------")
        # 调用 get_basic_financials 工具
        result = await client.call_tool('get_basic_financials', {'symbol': 'AAPL', 'selected_columns': ['eps', 'marketCapitalization']})
        print('调用结果:', result, "\n")  
        print("-------------------------------------------")
        # 调用 get_company_profile 工具
        profile = await client.call_tool('get_company_profile', {'symbol': 'AAPL'})
        print('公司简介:', profile, "\n") 
        print("-------------------------------------------")
        # 调用 get_company_news 工具
        news = await client.call_tool('get_company_news', {'symbol': 'AAPL', 'start_date': '2023-01-01', 'end_date': '2023-12-31'})
        print('公司新闻:', news, "\n")  

# 这是推荐的、唯一的程序入口点
if __name__ == '__main__':
    asyncio.run(main())
